# This code is hosted on http://code.google.com/p/lenthorp/
# Freely available for use in applications, but should NOT be modified
# Email all comments to lenthorpresearch@gmail.com


import PricingEngineMod
from PricingEngineMod import *
reload(PricingEngineMod)

import ExplicitFiniteDifferencesEngineMod
from ExplicitFiniteDifferencesEngineMod import *
reload(ExplicitFiniteDifferencesEngineMod)

from math import sqrt, log, exp

class BlackScholesHeatEquationExplicitFiniteDifferences(PricingEngine):

    def preComputations(self):
        p = dict(self.otherParams.p.items() + self.modelParams.p.items())
        strike = p['strike']
        maxTime = p['tenor'] * (p['vol'] ** 2.0) / 2.0
        variables = []
        maxStock = log(2*strike)
        stock = FiniteDifferenceVariable('stock', 0, maxStock)
        variables.append(stock)
        time = FiniteDifferenceVariable('time', 0, maxTime)
        variables.append(time)

##        partitions_stock = 100
##        # need: dt < 0.5 (dS)^2
##        partitions_time = int(1.1 * 2 * maxTime / ((float(maxStock)/partitions_stock) ** 2))

        rate = p['rate']
        vol = p['vol']
        
        partitions_time = int(100)# * maxTime) #10 per year!
        partitions_stock = int(100)#int(0.5 * 1 / (vol * sqrt(maxTime/partitions_time)))
        partitions = [partitions_stock, partitions_time]

        # PDE is u_t = u_XX
        #   u_t = 1/dt (u[i,j] - u[i-1,j])
        #   u_XX = 1/(dX)^2 (u[i-1,j+1] - 2 u[i-1,j] + u[i-1,j-1])
        #
        #   => u[i,j] = a_j u[i-1,j-1] + b_j u[i-1,j] + c_j u[i-1, j+1]
        #   where
        #       a_j = alpha
        #       b_j = 1 - 2 alpha
        #       c_j = alpha
        #   where
        #       alpha = dt / (dX)^2
        
        
        dt = maxTime / partitions_time

        relations = []
        #1:        
        relation = Relation()
        relation.variableNames = ['stock', 'time']
        relation.numberSteps = [-1, -1]
        relation.func = lambda x, y : x[1] / (x[0] ** 2.0)
        relations.append(relation)
        #2:        
        relation = Relation()
        relation.variableNames = ['stock', 'time']
        relation.numberSteps = [1, -1]
        relation.func = lambda x, y : x[1] / (x[0] ** 2.0)
        relations.append(relation)
        #3:        
        relation = Relation()
        relation.variableNames = ['stock', 'time']
        relation.numberSteps = [0, -1]
        relation.func = lambda x, y : 1.0 - 2.0 * x[1] / (x[0] ** 2.0) 
        relations.append(relation)

        initials = {}
        initials['stock'] = log(p['initial'])
        initials['time'] = maxTime

        a = rate / (vol ** 2) - 0.5
        bb = 2.0 * rate / (vol ** 2) + a**2.0

        self.fde = ExplicitFiniteDifferencesEngine(variables, partitions, 'time')
        # boundary conditions:
        #   t = Left: val = max(exp(X) - K,0) exp(-aX)
        #   s = Left: val = 0
        #   s = Right: val = (exp(X) - K,0) exp(-aX)

        #1        
        b = BoundaryCondition()
        b.conditionNames = ['time']
        b.LeftOrRight = 'Left'
        if p['type'] == 'call':
            b.funct = lambda x : 0.0 if exp(x[0]) < p['strike'] else (exp(x[0]) - p['strike']) * exp(a * x[0])
        elif p['type'] == 'put':
            b.funct = lambda x : 0.0 if exp(x[0]) > p['strike'] else (p['strike'] - exp(x[0])) * exp(a * x[0])
        else:
            raise Exception('unknown option type provided')
        self.fde.applyDirichletBoundaryCondition(b)

        #2
        b = BoundaryCondition()
        b.conditionNames = ['stock']
        b.LeftOrRight = 'Left'
        if p['type'] == 'call':
            b.funct = lambda x : 0.0
        elif p['type'] == 'put':
            b.funct = lambda x : (p['strike'] * exp(-2.0 * rate / (vol ** 2.0) * x[1])) * exp(a * x[0] + bb * x[1])
        self.fde.applyDirichletBoundaryCondition(b)

        #3  V[i,j] = 1/(1+rdt) (V[i+1,j] + rS dt) at the upper limit if a call (i.e. dS = 1), 0 if a put
        b = BoundaryCondition()
        b.conditionNames = ['stock']
        b.LeftOrRight = 'Right'
        if p['type'] == 'call':
            b.funct = lambda x : (exp(x[0]) - p['strike'] * exp(-2.0 * rate / (vol ** 2.0) * x[1])) * exp(a * x[0] + bb * x[1])
        elif p['type'] == 'put':
            b.funct = lambda x : 0.0
        self.fde.applyDirichletBoundaryCondition(b)
        
        self.fde.initialise(relations, initials)
        self.a = a
        self.b = bb
        self.initials = initials

        
    def getPrice(self):
        if (self.currentStatus > 0):
            op = self.otherParams.p
            mp = self.modelParams.p
            transformedPrice = self.fde.doCalculation()
            return transformedPrice * exp(-self.a * self.initials['stock'] - self.b * self.initials['time'])
            